#include "codegen.h"

CodeGenerator::CodeGenerator(ostream& output, bool opt) 
    : out(output), labelCounter(0), stackOffset(0), optimize(opt) {
    codeStream.str("");
    codeStream.clear();
}

string CodeGenerator::getGeneratedCode() const {
    return codeStream.str();
}

string CodeGenerator::newLabel() {
    return ".L" + to_string(labelCounter++);
}

void CodeGenerator::enterScope() {
    symbolTables.push(currentSymbols);
    currentSymbols = map<string, int>();
}

void CodeGenerator::exitScope() {
    if (!symbolTables.empty()) {
        currentSymbols = symbolTables.top();
        symbolTables.pop();
    }
}

void CodeGenerator::declareVariable(const string& name) {
    stackOffset -= 4; // 每个变量占4字节
    currentSymbols[name] = stackOffset;
}

int CodeGenerator::getVariableOffset(const string& name) {
    if (currentSymbols.find(name) != currentSymbols.end()) {
        return currentSymbols[name];
    }

    // 检查父作用域
    stack<map<string, int>> temp = symbolTables;
    while (!temp.empty()) {
        auto symbols = temp.top();
        temp.pop();
        if (symbols.find(name) != symbols.end()) {
            return symbols[name];
        }
    }

    cerr << "Error: Undeclared variable '" << name << "'" << endl;
    return 0;
}

void CodeGenerator::generate(CompUnitNode* node) {
    // 应用优化（如果启用）
    if (optimize) {
        applyOptimizations(node);
    }

    // 输出汇编头部
    codeStream << ".text\n";
    codeStream << ".globl main\n\n";

    // 生成所有函数
    for (auto& func : node->funcs) {
        if (auto funcDef = dynamic_cast<FuncDefNode*>(func.get())) {
            generateFunction(funcDef);
        }
    }

    out << codeStream.str();
}

// 新增优化方法实现
void CodeGenerator::applyOptimizations(CompUnitNode* node) {
    for (auto& func : node->funcs) {
        if (auto funcDef = dynamic_cast<FuncDefNode*>(func.get())) {
            optimizeFunction(funcDef);
        }
    }
}

void CodeGenerator::optimizeFunction(FuncDefNode* node) {
    // 常量折叠优化
    if (auto body = dynamic_cast<StmtNode*>(node->body.get())) {
        for (auto& stmt : body->children) {
            if (auto exprStmt = dynamic_cast<StmtNode*>(stmt.get())) {
                if (exprStmt->type == "ExprStmt" && !exprStmt->children.empty()) {
                    if (auto expr = dynamic_cast<ExprNode*>(exprStmt->children[0].get())) {
                        optimizeExpression(expr);
                    }
                }
            }
        }
    }
}

void CodeGenerator::optimizeExpression(ExprNode* node) {
    // 常量折叠：计算编译时可确定的表达式
    if (node->children.size() == 2) {
        auto left = dynamic_cast<NumberNode*>(node->children[0].get());
        auto right = dynamic_cast<NumberNode*>(node->children[1].get());
        
        if (left && right) {
            int lval = stoi(left->value);
            int rval = stoi(right->value);
            int result = 0;
            
            if (node->op == "+") result = lval + rval;
            else if (node->op == "-") result = lval - rval;
            else if (node->op == "*") result = lval * rval;
            else if (node->op == "/" && rval != 0) result = lval / rval;
            else if (node->op == "%" && rval != 0) result = lval % rval;
            
            // 替换为结果节点
            node->children.clear();
            node->op = to_string(result);
        }
    }
    
    // 递归优化子表达式
    for (auto& child : node->children) {
        if (auto expr = dynamic_cast<ExprNode*>(child.get())) {
            optimizeExpression(expr);
        }
    }
}


void CodeGenerator::generateFunction(FuncDefNode* node) {
    currentFunction = node->name;
    currentReturnType = node->retType;
    stackOffset = -8; // 初始偏移：fp(4) + ra(4)
    paramCount = 0;

    // 初始化作用域
    enterScope();
    symbolTables = stack<map<string, int>>();
    currentSymbols.clear();

    // 函数入口标签
    codeStream << node->name << ":\n";

    // 保存帧指针和返回地址
    codeStream << "    addi sp, sp, -8\n";
    codeStream << "    sw ra, 4(sp)\n";
    codeStream << "    sw fp, 0(sp)\n";
    codeStream << "    mv fp, sp\n";

    // 处理函数参数
    for (auto& param : node->params) {
        if (auto expr = dynamic_cast<ExprNode*>(param.get())) {
            if (expr->op == "Param" && !expr->children.empty()) {
                if (auto ident = dynamic_cast<ExprNode*>(expr->children[0].get())) {
                    string paramName = ident->op;
                    declareVariable(paramName);
                    codeStream << "    sw a" << paramCount << ", " << getVariableOffset(paramName) << "(fp)\n";
                    paramCount++;
                }
            }
        }
    }

    // 生成函数体
    if (auto body = dynamic_cast<StmtNode*>(node->body.get())) {
        generateStatement(body);
    }

    // 处理void函数无return的情况
    if (currentReturnType == "void") {
        codeStream << "    lw ra, 4(sp)\n";
        codeStream << "    lw fp, 0(sp)\n";
        codeStream << "    addi sp, sp, 8\n";
        codeStream << "    ret\n";
    }

    // 函数结束标签
    codeStream << ".end_" << node->name << ":\n\n";
    exitScope();
}

void CodeGenerator::generateStatement(StmtNode* node) {
    if (node->type == "Block") {
        enterScope();
        for (auto& child : node->children) {
            if (auto stmt = dynamic_cast<StmtNode*>(child.get())) {
                generateStatement(stmt);
            }
        }
        exitScope();
    }
    else if (node->type == "DeclareStmt") {
        if (!node->children.empty()) {
            if (auto ident = dynamic_cast<ExprNode*>(node->children[0].get())) {
                string varName = ident->op;
                declareVariable(varName);

                // 处理初始化表达式
                if (node->children.size() > 1) {
                    generateExpression(node->children[1].get());
                    codeStream << "    sw t0, " << getVariableOffset(varName) << "(fp)\n";
                }
            }
        }
    }
    else if (node->type == "AssignStmt") {
        if (node->children.size() >= 2) {
            if (auto ident = dynamic_cast<ExprNode*>(node->children[0].get())) {
                string varName = ident->op;
                generateExpression(node->children[1].get());
                codeStream << "    sw t0, " << getVariableOffset(varName) << "(fp)\n";
            }
        }
    }
    else if (node->type == "ReturnStmt") {
        if (!node->children.empty()) {
            generateExpression(node->children[0].get());
            codeStream << "    mv a0, t0\n";
        }
        codeStream << "    lw ra, 4(sp)\n";
        codeStream << "    lw fp, 0(sp)\n";
        codeStream << "    addi sp, sp, 8\n";
        codeStream << "    ret\n";
    }
    else if (node->type == "IfStmt") {
        string elseLabel = newLabel();
        string endLabel = newLabel();

        // 条件表达式
        generateExpression(node->children[0].get());
        codeStream << "    beqz t0, " << elseLabel << "\n";

        // then分支
        if (auto stmt = dynamic_cast<StmtNode*>(node->children[1].get())) {
            generateStatement(stmt);
        }
        codeStream << "    j " << endLabel << "\n";

        // else分支
        codeStream << elseLabel << ":\n";
        if (node->children.size() > 2) {
            if (auto stmt = dynamic_cast<StmtNode*>(node->children[2].get())) {
                generateStatement(stmt);
            }
        }

        codeStream << endLabel << ":\n";
    }
    else if (node->type == "WhileStmt") {
        string startLabel = newLabel();
        string endLabel = newLabel();

        // 保存循环标签
        loopStartLabels.push(startLabel);
        loopEndLabels.push(endLabel);

        codeStream << startLabel << ":\n";
        // 条件判断
        generateExpression(node->children[0].get());
        codeStream << "    beqz t0, " << endLabel << "\n";

        // 循环体
        if (auto stmt = dynamic_cast<StmtNode*>(node->children[1].get())) {
            generateStatement(stmt);
        }
        codeStream << "    j " << startLabel << "\n";
        codeStream << endLabel << ":\n";

        loopStartLabels.pop();
        loopEndLabels.pop();
    }
    else if (node->type == "BreakStmt") {
        if (!loopEndLabels.empty()) {
            codeStream << "    j " << loopEndLabels.top() << "\n";
        }
    }
    else if (node->type == "ContinueStmt") {
        if (!loopStartLabels.empty()) {
            codeStream << "    j " << loopStartLabels.top() << "\n";
        }
    }
    else if (node->type == "ExprStmt") {
        if (!node->children.empty()) {
            generateExpression(node->children[0].get());
        }
    }
}

void CodeGenerator::generateExpression(ASTNode* node) {
    // 处理数字节点
    if (auto numNode = dynamic_cast<NumberNode*>(node)) {
        codeStream << "    li t0, " << numNode->value << "\n";
        return;
    }

    // 处理表达式节点
    if (auto exprNode = dynamic_cast<ExprNode*>(node)) {
        // 处理标识符（变量引用）
        if (exprNode->children.empty()) {
            int offset = getVariableOffset(exprNode->op);
            codeStream << "    lw t0, " << offset << "(fp)\n";
            return;
        }

        // 处理函数调用
        if (exprNode->op == "Primary" && exprNode->children.size() > 1) {
            if (auto funcNameNode = dynamic_cast<ExprNode*>(exprNode->children[0].get())) {
                string funcName = funcNameNode->op;
                int argCount = 0;

                // 保存寄存器（保护t1并保证栈对齐）
                codeStream << "    addi sp, sp, -8\n";
                codeStream << "    sw t1, 4(sp)\n";
                codeStream << "    sw t2, 0(sp)\n";

                // 设置参数
                for (size_t i = 1; i < exprNode->children.size(); i++) {
                    generateExpression(exprNode->children[i].get());
                    codeStream << "    mv a" << argCount << ", t0\n";
                    argCount++;
                }

                codeStream << "    call " << funcName << "\n";
                codeStream << "    mv t0, a0\n";  // 返回值存入t0

                // 恢复寄存器
                codeStream << "    lw t2, 0(sp)\n";
                codeStream << "    lw t1, 4(sp)\n";
                codeStream << "    addi sp, sp, 8\n";
                return;
            }
        }

        // 处理一元运算符
        if (exprNode->children.size() == 1) {
            generateExpression(exprNode->children[0].get());
            if (exprNode->op == "!") {
                codeStream << "    seqz t0, t0\n";
            }
            else if (exprNode->op == "-") {
                codeStream << "    neg t0, t0\n";
            }
            return;
        }

        // 处理二元运算符
        if (exprNode->children.size() == 2) {
            // 修复短路求值逻辑
            if (exprNode->op == "&&") {
                string falseLabel = newLabel();
                string endLabel = newLabel();

                // 先计算左操作数
                generateExpression(exprNode->children[0].get());
                codeStream << "    beqz t0, " << falseLabel << "\n"; // 左操作数为假则跳转
                
                // 再计算右操作数
                generateExpression(exprNode->children[1].get());
                codeStream << "    beqz t0, " << falseLabel << "\n"; // 右操作数为假则跳转
                
                codeStream << "    li t0, 1\n";
                codeStream << "    j " << endLabel << "\n";
                
                codeStream << falseLabel << ":\n";
                codeStream << "    li t0, 0\n";
                
                codeStream << endLabel << ":\n";
                return;
            }
            else if (exprNode->op == "||") {
                string trueLabel = newLabel();
                string endLabel = newLabel();

                // 先计算左操作数
                generateExpression(exprNode->children[0].get());
                codeStream << "    bnez t0, " << trueLabel << "\n"; // 左操作数为真则跳转
                
                // 再计算右操作数
                generateExpression(exprNode->children[1].get());
                codeStream << "    bnez t0, " << trueLabel << "\n"; // 右操作数为真则跳转
                
                codeStream << "    li t0, 0\n";
                codeStream << "    j " << endLabel << "\n";
                
                codeStream << trueLabel << ":\n";
                codeStream << "    li t0, 1\n";
                
                codeStream << endLabel << ":\n";
                return;
            }

            // 修复运算符优先级问题
            // 保存左操作数到栈上
            generateExpression(exprNode->children[0].get());
            codeStream << "    addi sp, sp, -4\n";
            codeStream << "    sw t0, 0(sp)\n";
            
            // 计算右操作数
            generateExpression(exprNode->children[1].get());
            
            // 恢复左操作数
            codeStream << "    lw t1, 0(sp)\n";
            codeStream << "    addi sp, sp, 4\n";
            
            // 执行运算
            if (exprNode->op == "+") codeStream << "    add t0, t1, t0\n";
            else if (exprNode->op == "-") codeStream << "    sub t0, t1, t0\n";
            else if (exprNode->op == "*") codeStream << "    mul t0, t1, t0\n";
            else if (exprNode->op == "/") codeStream << "    div t0, t1, t0\n";
            else if (exprNode->op == "%") codeStream << "    rem t0, t1, t0\n";
            else if (exprNode->op == "<") codeStream << "    slt t0, t1, t0\n";
            else if (exprNode->op == ">") codeStream << "    slt t0, t0, t1\n";
            else if (exprNode->op == "<=") {
                codeStream << "    slt t0, t0, t1\n";
                codeStream << "    xori t0, t0, 1\n";
            }
            else if (exprNode->op == ">=") {
                codeStream << "    slt t0, t1, t0\n";
                codeStream << "    xori t0, t0, 1\n";
            }
            else if (exprNode->op == "==") {
                codeStream << "    xor t0, t1, t0\n";
                codeStream << "    seqz t0, t0\n";
            }
            else if (exprNode->op == "!=") {
                codeStream << "    xor t0, t1, t0\n";
                codeStream << "    snez t0, t0\n";
            }
            return;
        }
    }

    cerr << "Error: Unhandled expression type" << endl;
}